package info.wangyuan.agent.service.curd.impl;

import cn.hutool.core.codec.Base64;
import com.baomidou.mybatisplus.core.toolkit.Wrappers;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import dev.langchain4j.data.message.ImageContent;
import dev.langchain4j.data.message.SystemMessage;
import dev.langchain4j.data.message.TextContent;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;
import info.wangyuan.agent.common.Constant.ModelConstant;
import info.wangyuan.agent.common.Constant.PromptTemplate;
import info.wangyuan.agent.common.Enum.ModelType;
import info.wangyuan.agent.config.ModelFactory;
import info.wangyuan.agent.entity.po.ModelConfig;
import info.wangyuan.agent.exception.BusinessException;
import info.wangyuan.agent.mapper.ModelConfigMapper;
import info.wangyuan.agent.service.curd.ModelConfigService;
import info.wangyuan.agent.service.rag.RerankDomainService;
import info.wangyuan.agent.utils.TikaFileTypeDetector;
import lombok.extern.slf4j.Slf4j;
import org.mockito.Mockito;
import org.springframework.stereotype.Service;

import java.util.List;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.stream.IntStream;

/**
 * 模型配置服务实现类
 *
 * @author Albert
 * @since 2025-07-11 15:23:42
 */
@Slf4j
@Service
public class ModelConfigServiceImpl extends ServiceImpl<ModelConfigMapper, ModelConfig> implements ModelConfigService {

    private final ModelFactory modelFactory;
    private final ModelConfigMapper modelConfigMapper;
    private final RerankDomainService rerankDomainService;

    public ModelConfigServiceImpl(ModelFactory modelFactory, ModelConfigMapper modelConfigMapper, RerankDomainService rerankDomainService) {
        this.modelFactory = modelFactory;
        this.modelConfigMapper = modelConfigMapper;
        this.rerankDomainService = rerankDomainService;
    }

    /**
     * 添加模型
     */
    public void canAdd(ModelConfig modelConfig) {
        boolean exists = lambdaQuery().eq(ModelConfig::getName, modelConfig.getName()).eq(ModelConfig::getType, modelConfig.getType()).eq(ModelConfig::getProviderId, modelConfig.getProviderId()).exists();
        if (exists) {
            throw new BusinessException("模型已存在");
        }
        // 测试模型可用性
        modelTest(modelConfig);
        modelConfigMapper.insert(modelConfig);
    }

    /**
     * 获取模型列表
     */
    public List<ModelConfig> getModelList(String type) {
        if (type == null) return modelConfigMapper.selectList(null);
        return modelConfigMapper.selectList(Wrappers.lambdaQuery(ModelConfig.class).eq(ModelConfig::getType, type));
    }

    /**
     * 删除模型
     */
    public void canDelete(Integer id) {
        ModelConfig modelConfig = getById(id);

        validateModelExistence(modelConfig);
        validateModelNotActive(modelConfig);
        validateModelTypeCount(modelConfig);
        validateModelNotUsed(modelConfig);

        modelConfigMapper.deleteById(modelConfig);
    }

    /**
     * 校验模型是否存在
     */
    private void validateModelExistence(ModelConfig modelConfig) {
        if (modelConfig == null) {
            throw new BusinessException("模型不存在");
        }
    }

    /**
     * 校验模型是否正在使用
     */
    private void validateModelNotActive(ModelConfig modelConfig) {
        if (modelConfig.getStatus().equals(ModelConstant.IS_ACTIVE)) {
            throw new BusinessException("模型正在被使用，无法删除");
        }
    }

    /**
     * 校验模型类型是否至少保留一个
     */
    private void validateModelTypeCount(ModelConfig modelConfig) {
        long sameTypeCount = modelConfigMapper.selectCount(Wrappers.lambdaQuery(ModelConfig.class).eq(ModelConfig::getType, modelConfig.getType()));
        if (sameTypeCount <= 1) {
            throw new BusinessException("该类型至少保留一个模型，无法删除");
        }
    }

    /**
     * 检查模型是否被使用
     */
    private void validateModelNotUsed(ModelConfig modelConfig) {
        Boolean flag = modelConfigMapper.selectOne(Wrappers.lambdaQuery(ModelConfig.class).eq(ModelConfig::getId, modelConfig.getId())).getStatus();

        if (Boolean.FALSE.equals(flag)) return;
        throw new BusinessException("该模型正在使用，无法删除");
    }

    /**
     * 修改模型
     */
    public void canUpdate(ModelConfig modelConfig) {
        ModelConfig dbModelConfig = modelConfigMapper.selectById(modelConfig.getId());
        long sameTypeCount = modelConfigMapper.selectCount(Wrappers.lambdaQuery(ModelConfig.class).eq(ModelConfig::getType, dbModelConfig.getType()));

        if (sameTypeCount == 1 && !modelConfig.getType().equals(dbModelConfig.getType())) {
            throw new BusinessException("该类型至少保留一个模型，无法修改");
        }
        modelConfigMapper.updateById(modelConfig);
    }

    private static final Pattern BOX_PATTERN = Pattern.compile("<\\|begin_of_box\\|>(.*?)<\\|end_of_box\\|>");

    /**
     * 模型测试入口
     */
    public void modelTest(ModelConfig modelConfig) {
        String type = modelConfig.getType();

        if (type.equals(ModelType.LLM.getValue())) {
            testChatModel();
        } else if (type.equals(ModelType.TEXT_EMBEDDING.getValue())) {
            testEmbeddingModel();
        } else if (type.equals(ModelType.RERANK.getValue())) {
            testRerankModel();
        } else if (type.equals(ModelType.OCR.getValue())) {
            testOcrModel();
        } else {
            throw new BusinessException("该模型无法正常使用，请检查配置");
        }
    }

    /**
     * 对话模型测试
     */
    private void testChatModel() {
        runWithLog("[对话模型测试]", () -> {
            String result = modelFactory.getCustomChatModel()
                    .chat(new SystemMessage("测试模型输出"), new UserMessage("回复一个字"))
                    .toString();
            log.info("[对话模型测试结果] | {}", result);
            if (result.isBlank()) {
                throw new BusinessException("该模型无法正常使用，请检查配置");
            }
        });
    }

    /**
     * 嵌入模型测试
     */
    private void testEmbeddingModel() {
        runWithLog("[嵌入模型测试]", () -> {
            int dimension = modelFactory.getCustomEmbeddingModel().embed("测试").content().dimension();
            log.info("[嵌入模型测试结果] | 输出维度：{}", dimension);
            if (dimension != ModelConstant.CUSTOM_DIMENSIONS) {
                throw new BusinessException(
                        String.format("嵌入模型仅支持输出维度：%d | 当前输出维度：%d", ModelConstant.CUSTOM_DIMENSIONS, dimension));
            }
        });
    }

    /**
     * 重排序模型测试
     */
    private void testRerankModel() {
        runWithLog("[重排序模型测试]", () -> {
            List<TextSegment> segments = List.of(TextSegment.from("A"), TextSegment.from("B"));
            List<EmbeddingMatch<TextSegment>> matchList = IntStream.range(0, segments.size())
                    .mapToObj(i -> new EmbeddingMatch<>(1.0, "test-id-" + i, null, segments.get(i)))
                    .toList();

            @SuppressWarnings("unchecked")
            EmbeddingSearchResult<TextSegment> searchResult = Mockito.mock(EmbeddingSearchResult.class);
            Mockito.when(searchResult.matches()).thenReturn(matchList);

            List<EmbeddingMatch<TextSegment>> rerankResult = rerankDomainService.rerankDocument(searchResult, "Test query");

            log.info("[重排序模型测试结果] | {}", rerankResult.stream().map(m -> m.embedded().text()).toList());
            if (rerankResult.size() != segments.size()) {
                throw new BusinessException("该模型无法正常使用，请检查配置");
            }
        });
    }

    /**
     * OCR 模型测试
     */
    private void testOcrModel() {
        runWithLog("[图像识别模型测试]", () -> {
            final UserMessage userMessage = new UserMessage(
                    ImageContent.from(ModelConstant.TEST_OCR_IMG, TikaFileTypeDetector.detectFileType(Base64.decode(ModelConstant.TEST_OCR_IMG))),
                    TextContent.from(PromptTemplate.TEST_OCR_PROMPT)
            );

            String vlmResult = modelFactory.getCustomOcrModel().chat(userMessage).aiMessage().text().trim();
            log.info("[图像识别模型测试结果] | {}", vlmResult);

            Matcher matcher = BOX_PATTERN.matcher(vlmResult);
            if (matcher.find()) {
                String content = matcher.group(1).trim();
                if (!content.equals("6")) {
                    throw new BusinessException("该模型图像识别异常，请检查配置");
                }
            } else {
                throw new BusinessException("该模型无法正常使用，请检查配置");
            }
        });
    }

    /**
     * 统一日志与异常捕获
     */
    private void runWithLog(String title, Runnable action) {
        try {
            log.info("{}开始", title);
            action.run();
            log.info("{}结束", title);
        } catch (Exception e) {
            log.error("{}失败：{}", title, e.getMessage());
            throw new BusinessException("该模型无法正常使用，请检查配置");
        }
    }

}
